""" This script builds the value functions

"""
if __name__ == '__main__':
    #%% IMPORT --------------------------------------------------------------------------
    import pandas as pd
    import numpy as np
    import sys
    sys.path.append('./')

    from src.models_new import value_search, construct_data

    #%% SET THE PARAMETERS --------------------------------------------------------------
    g_grid = 10
    n_grid = 5
    beta_by_time = {'month': 0.99, 'fortnight': 0.995}
    sim_length = 1200
    seeds = list(range(200))
    options = {
        'threads_per_worker': 1,
        'n_workers': 8
    }

    #%% READ IN THE PARAMETERS ----------------------------------------------------------
    for i in ['month', 'fortnight']:
        if i == 'month':
            fortnight = False
        elif i == 'fortnight':
            fortnight = True
        df_ols = pd.read_csv(f'./models/first_stage/first_stage_ols_{i}.csv', index_col=[0, 1])
        df_mnl = pd.read_csv(f'./models/first_stage/first_stage_mnl_{i}.csv', index_col=[0, 1])
        df_extensions = pd.read_csv(f'./models/first_stage/first_stage_extensions_{i}.csv', index_col=[0, 1])

        r = pd.read_csv(f'./models/first_stage/r_{i}.csv', index_col=[0])
        const = pd.read_csv(f'./models/first_stage/const_{i}.csv', index_col=[0])
        sigma = pd.read_csv(f'./models/first_stage/sigma_{i}.csv', index_col=[0])

        #%% BUILD THE FIRST STAGE FUNCTIONS ---------------------------------------------
        first_stage_by_spec = dict()
        for spec in ['low', 'mid', 'high']:
            first_stage_by_spec[spec] = {
                'price_new': construct_data.build_ols_predictor(
                    coefs=df_ols.loc[(spec, 'day_rate')].values),
                'price_extend': construct_data.build_price_extension_predictor(
                    coefs=df_extensions.loc[(spec, 'price')].values),
                'prob_match': construct_data.build_mnl_predictor(
                    coefs=df_mnl.loc[spec].values),
                'prob_extend': construct_data.build_prob_extension_predictor(
                    coefs=df_extensions.loc[(spec, 'prob')].values),
                'const': const.values.T[0],
                'r': r.values,
                'sigma': sigma.values[0][0]
            }

        #%% BUILD THE VALUE FUNCTIONS ---------------------------------------------------
        search_grid_by_spec = dict()
        search_value_by_spec = dict()
        search_value_with_node_by_spec = dict()
        for spec in ['low', 'mid', 'high']:
            (
                search_grid_by_spec[spec],
                search_value_by_spec[spec],
                search_value_with_node_by_spec[spec]
            ) = value_search.build_value_functions(
                g_grid, n_grid, first_stage_by_spec[spec], beta_by_time[i],
                sim_length, seeds, fortnight, options
            )

        #%% SAVE THE VALUE FUNCTIONS ----------------------------------------------------
        for spec in ['low', 'mid', 'high']:
            np.save(
                f'./models/value_search/search_grid_{spec}_{i}.npy',
                search_grid_by_spec[spec]
            )
            np.save(
                f'./models/value_search/search_value_{spec}_{i}.npy',
                search_value_by_spec[spec]
            )
            pd.Series(search_value_with_node_by_spec[spec]).to_csv(
                f'./models/value_search/search_value_with_node_{spec}_{i}.csv'
            )
